package com.gee.spring.boot.starter.opengemini.config;

import com.gee.spring.boot.starter.opengemini.core.domain.MeasurementInfo;
import com.gee.spring.boot.starter.opengemini.core.domain.Page;
import com.gee.spring.boot.starter.opengemini.core.domain.PointWrapper;
import com.gee.spring.boot.starter.opengemini.core.query.QueryWrapper;
import com.gee.spring.boot.starter.opengemini.core.util.ExceptionUtil;
import com.gee.spring.boot.starter.opengemini.core.util.OpengeminiUtil;
import io.opengemini.client.api.Point;
import io.opengemini.client.api.Query;
import io.opengemini.client.api.QueryResult;
import io.opengemini.client.jdk.OpenGeminiJdkClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static com.gee.spring.boot.starter.opengemini.core.util.OpengeminiUtil.*;

/**
 * @author gepengjun
 * @since 2024/4/30 9:55
 */
public class OpengeminiTemplate {

    private static final Logger logger = LoggerFactory.getLogger(OpengeminiTemplate.class);

    private final OpenGeminiJdkClient openGeminiJdkClient;

    private final OpengeminiProperties opengeminiProperties;

    private final int defaultBatchInsertSize;

    public OpengeminiTemplate(OpenGeminiJdkClient openGeminiJdkClient, OpengeminiProperties opengeminiProperties) {
        this.openGeminiJdkClient = openGeminiJdkClient;
        this.opengeminiProperties = opengeminiProperties;
        this.defaultBatchInsertSize = opengeminiProperties.getDefaultBatchInsertSize();
    }

    public QueryResult execute(Query query){
        try {
            logger.info("execute query : {}, database: {}, retentionPolicy: {}", query.getCommand(),
                    query.getDatabase(), query.getRetentionPolicy());
            return openGeminiJdkClient.query(query).get();
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    public <T> List<T> list(QueryWrapper<T> queryWrapper) {
        return OpengeminiUtil.toPOJO(execute(getQueryByLambda(queryWrapper)), queryWrapper.getEntityClass(), queryWrapper.getFuncTime());
    }

    public <T> Page<T> page(QueryWrapper<T> queryWrapper, Page<T> page) {
        QueryResult countResult = execute(getCountQueryByLambda(queryWrapper));
        long count = getCount(countResult);
        if (count > 0){
            page.setTotal(count);
            queryWrapper.limit(page.getPageSize());
            if (page.getCurrent() > 1){
                queryWrapper.offset((page.getCurrent() - 1 ) * page.getPageSize());
            }
            page.setDataList(list(queryWrapper));
        }
        return page;
    }

    public void batchInsert(List<?> entityList) {
        batchInsert(entityList, defaultBatchInsertSize);
    }

    public void batchInsert(List<?> entityList, int batchSize) {
        if (entityList.isEmpty()){
            return;
        }
        if (batchSize <= 0){
            throw ExceptionUtil.re("batch size should be greater than 0");
        }
        List<PointWrapper> pointWrapperList = new ArrayList<>();
        for (Object object : entityList) {
            pointWrapperList.add(OpengeminiUtil.toPointWrapper(object));
        }
        Map<String, List<PointWrapper>> groupByDatabase = pointWrapperList.stream().collect(Collectors.groupingBy(PointWrapper::getDatabase));
        groupByDatabase.forEach((database, wrapperList) -> {
            int size = wrapperList.size();
            List<Point> pointList = wrapperList.stream().map(PointWrapper::getPoint).collect(Collectors.toList());
            if (size < batchSize){
                openGeminiJdkClient.writeBatch(database, pointList);
            }else {
                int batchMulti = size / batchSize;
                int batchCount = size % batchSize == 0 ? batchMulti : batchMulti + 1;
                IntStream.range(0, batchCount)
                        .mapToObj(batchNum -> pointList.subList(batchNum * batchSize, Math.min(size, (batchNum + 1) * batchSize)))
                        .forEach(list -> openGeminiJdkClient.writeBatch(database, list));
            }
        });
    }

    public void insert(Object entity) {
        Assert.notNull(entity, "entity should not be null");
        MeasurementInfo measurementInfo = getMeasurementInfo(entity.getClass());
        String database = measurementInfo.getDatabase(entity);

        openGeminiJdkClient.write(database, toPointWrapper(entity).getPoint());
    }

    private static <T> Query getCountQueryByLambda(QueryWrapper<T> queryWrapper){
        MeasurementInfo measurementInfo = OpengeminiUtil.getMeasurementInfo(queryWrapper.getEntityClass());
        return new Query(queryWrapper.toCountSql(), queryWrapper.getDatabase(), measurementInfo.getRetentionPolicy());
    }

    private static <T> Query getQueryByLambda(QueryWrapper<T> queryWrapper){
        MeasurementInfo measurementInfo = OpengeminiUtil.getMeasurementInfo(queryWrapper.getEntityClass());
        return new Query(queryWrapper.toSqlStr(), queryWrapper.getDatabase(), measurementInfo.getRetentionPolicy());
    }

}
